from langchain.tools import tool
import os
from langchain.chat_models import init_chat_model
from dotenv import load_dotenv
from langchain_classic.agents import create_tool_calling_agent, AgentExecutor
from langchain_core.messages import HumanMessage, AIMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder

# 加载.env文件

load_dotenv()

# 第一步：定义加法、减法工具
# 通过@tool装饰器分别定义两个工具，关键是在文档字符串中明确区分功能，让模型能精准识别适用场景。

# 加法工具
@tool
def add_tool(num1: float, num2: float) -> str:
    """
    用于执行两个数字的加法运算。
    当用户需要计算“A加B”“A加上B”“A与B的和”“A和B相加”等场景时使用。

    Args:
        num1: 加法运算中的第一个数字（整数或小数均可）
        num2: 加法运算中的第二个数字（整数或小数均可）
    """
    result = num1 + num2
    return f"加法运算结果：{num1} + {num2} = {result}"

# 减法工具
@tool
def subtract_tool(num1: float, num2: float) -> str:
    """
    用于执行两个数字的减法运算。
    当用户需要计算“A减B”“A减去B”“A与B的差”“A比B多多少”“A减去B等于多少”等场景时使用。

    Args:
        num1: 减法运算中的被减数（整数或小数均可）
        num2: 减法运算中的减数（整数或小数均可）
    """
    result = num1 - num2
    return f"减法运算结果：{num1} - {num2} = {result}"


# 第二步：搭建 Agent 让模型驱动工具选择
# 将两个工具传入 Agent，结合大语言模型（LLM），让模型根据用户输入的问题，自主判断该调用加法还是减法工具。
# 1. 初始化模型
model = init_chat_model(
    model='deepseek-reasoner', # deepseek-chat表示调用DeepSeek-v3模型，deepseek-reasoner表示调用DeepSeek-R1模型，
    model_provider='deepseek',# 模型提供商写deepseek
    # api_key="sk-476b894aa2d74474b77d4dfa91633f72", #你注册的deepseek api_key
    api_key=os.environ.get("DEEPSEEK_API_KEY"),
)
tools = [add_tool, subtract_tool]
# 2. 构建带对话记忆的 Prompt
# 关键是加入MessagesPlaceholder存储对话历史，同时保留agent_scratchpad记录工具调用流程。

# 构建Prompt：包含对话历史、当前输入、工具调用记录
prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一个能进行多轮对话的数学助手，需基于完整对话上下文回应。"
               "如果需要计算，务必根据需求选择add_tool或subtract_tool，严格按参数传参。"
               "不需要计算时，直接用自然语言回复即可。"),
    # 对话历史占位符：存储所有过往消息
    MessagesPlaceholder(variable_name="chat_history"),
    # 当前用户输入
    ("human", "{input}"),
    # 工具调用记录占位符：存储工具调用请求和结果
    ("placeholder", "{agent_scratchpad}")
])

# 3. 创建 Agent 执行器 & 实现对话循环
# 通过循环接收用户输入，每次将新消息追加到chat_history，实现多轮对话上下文管理。

# 创建工具调用Agent
agent = create_tool_calling_agent(model, tools, prompt)

# 创建Agent执行器（开启verbose查看思考过程）
agent_executor = AgentExecutor(
    agent=agent,
    tools=tools,
    verbose=True,
    handle_parsing_errors=True
)

# 初始化对话历史（空列表，将逐步追加消息）
chat_history = []

print("数学助手已启动，输入'退出'结束对话\n")

# 对话循环
while True:
    # 获取用户输入
    user_input = input("你：")
    if user_input == "退出":
        print("助手：再见！")
        break

    # 调用Agent，传入当前输入和对话历史
    response = agent_executor.invoke({
        "input": user_input,
        "chat_history": chat_history  # 传入历史上下文
    })

    # 提取助手回复内容
    assistant_response = response["output"]
    print(f"助手：{assistant_response}")

    # 更新对话历史：追加本轮用户消息和助手消息
    chat_history.append(HumanMessage(content=user_input))
    chat_history.append(AIMessage(content=assistant_response))